# 54/100 图论-实现Trie(前缀树)
# leetcode第208题: https://leetcode.cn/problems/implement-trie-prefix-tree/description/?envType=study-plan-v2&envId=top-100-liked
# Date: 2025/1/13
"""Trie(也叫 前缀树 字典树 单词查找树)
字典树是一种常见的数据结构, 主要用于高效地存储和检索字符串数据集中的键. 这一数据结构有相当多的应用场景, 如自动补全和拼写检查等等.
如常见的搜索引擎提供的自动补全和常用文字编辑器的拼写检查等.

如创建一个关键字集合["a", "app", "apple", "to", "tea", "ted", "ten", "in", "inn"]
那么Trie树就可以表示为:
Root:
    a: *
        p:
            p: *
                l:
                    e: *
    i:
        n: *
        nn: *
    t:
        o: *
        e:
            a: *
            d: *
            n: *
* 表示该节点是一个单词的结束。
"""
import sys

from leetcode import bds


class Trie:
    def __init__(self):
        # 初始化子节点数组和结束标志
        self.children = [None] * 26
        self.is_end = False

    def insert(self, word: str) -> None:
        node = self
        for ch in word:
            index = ord(ch) - ord('a')
            if not node.children[index]:
                node.children[index] = Trie()
            node = node.children[index]
        node.is_end = True

    def search(self, word: str) -> bool:
        node = self._search_prefix(word)
        return node is not None and node.is_end

    def startsWith(self, prefix: str) -> bool:
        return self._search_prefix(prefix) is not None

    def _search_prefix(self, prefix: str) -> 'Trie':
        node = self
        for ch in prefix:
            index = ord(ch) - ord('a')
            if not node.children[index]:
                return None
            node = node.children[index]
        return node

    def __str__(self):
        return str(self.children)


class TrieOpt:
    """
    实现 Trie (前缀树) 数据结构。
    功能：
    - insert(word): 插入一个字符串到 Trie 中。
    - search(word): 检查 Trie 中是否完全包含某个字符串（非前缀形式）。
    - startsWith(prefix): 检查 Trie 中是否有字符串以该前缀开始。

    区别于Trie, 这里使用的是递归的dict来存储树结构, 虽然运行的效率更高, 但是需要消耗大量的内存
    例子:
    如 "apple" 就是:
    {'a': {'p': {'p': {'l': {'e': {'*': ''}}}}}}
    继续添加"app" 就是:
    {'a':
        {'p': {'p': {'l': {'e': {'*': ''}}, '*': ''}}}}
    """

    def __init__(self):
        self.trie = {}

    def insert(self, word: str) -> None:
        """
        插入一个单词到 Trie 中。
        从 Trie 的根开始，逐步添加每个字符到 Trie 结构中，如果字符不存在，则创建新的字典节点。
        """
        cur = self.trie
        for c in word:
            if c not in cur:
                cur[c] = {}
            cur = cur[c]

        cur['*'] = ''

    def search(self, word: str) -> bool:
        """
        搜索 Trie 中是否存在一个完整的单词。
        按照单词的每个字符逐层深入 Trie，如果在任何点字符路径中断，则该单词不存在。
        如果所有字符都正确匹配，最后检查是否有结束标记 '*'。
        """
        cur = self.trie
        for c in word:
            if c in cur:
                cur = cur[c]
            else:
                return False
        if '*' in cur:
            return True
        else:
            return False

    def startsWith(self, prefix: str) -> bool:
        """
        检查 Trie 中是否有单词以指定的前缀开始。
        这类似于 search 方法，但不需要检查单词结束标记 '*'，只需要确认前缀路径存在。
        """
        cur = self.trie
        for c in prefix:
            if c in cur:
                cur = cur[c]
            else:
                return False
        return True

    def __str__(self):
        return str(self.trie)


if __name__ == '__main__':
    trie = Trie()
    trie.insert("apple")
    print(trie.search("apple"))  # True
    print(trie.search("app"))  # False
    print(trie.startsWith("app"))  # True
    trie.insert("app")
    trie.insert("a")
    trie.insert("to")
    trie.insert("tea")
    trie.insert("ted")
    trie.insert("ten")
    trie.insert("i")
    trie.insert("in")
    trie.insert("inn")
    print(trie.search("app"))  # True
    print(bds.total_size(trie.children) + bds.total_size(trie.is_end))  # 448

    print("=" * 25)

    trie_o = TrieOpt()
    trie_o.insert("apple")
    print(trie_o.search("apple"))  # True
    print(trie_o.search("app"))  # False
    print(trie_o.startsWith("app"))  # True
    trie_o.insert("app")
    trie_o.insert("a")
    trie_o.insert("to")
    trie_o.insert("tea")
    trie_o.insert("ted")
    trie_o.insert("ten")
    trie_o.insert("i")
    trie_o.insert("in")
    trie_o.insert("inn")
    print(trie.search("app"))  # True
    print(bds.total_size(trie_o.trie))  # 4029
